{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Building the Transformer from Scratch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook, we'll be implementing the famous Transformer architecture from scratch.\n",
    "\n",
    "The code is based off of the following repos/blog posts:\n",
    "\n",
    "- [attention-is-all-you-need-pytorch](https://github.com/jadore801120/attention-is-all-you-need-pytorch)\n",
    "- [pytorch-pretrained-BERT](https://github.com/huggingface/pytorch-pretrained-BERT)\n",
    "- [The Annotated Transformer](http://nlp.seas.harvard.edu/2018/04/03/attention.html) \n",
    "\n",
    "Thanks so much to their authors!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One of the keys to understanding how any model works is understanding how the shapes of the tensors change during the processing of each part. We'll be using the logging module to output debugging information to help our understanding."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "logger = logging.getLogger(\"tensor_shapes\")\n",
    "handler = logging.StreamHandler()\n",
    "formatter = logging.Formatter(\n",
    "        '%(message)s')\n",
    "handler.setFormatter(formatter)\n",
    "logger.addHandler(handler)\n",
    "# if you want the model to continuously print tensor shapes, set to DEBUG!\n",
    "logger.setLevel(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import inspect\n",
    "def getclass():\n",
    "    stack = inspect.stack()\n",
    "    return stack[3][0].f_locals[\"self\"].__class__\n",
    "\n",
    "# A helper function to check how tensor sizes change\n",
    "def log_size(tsr: torch.Tensor, name: str):\n",
    "    cls = getclass()\n",
    "    logger.log(level=cls.level, msg=f\"[{cls.__name__}] {name} size={tsr.shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll use logging levels to control the modules we receive output from. The lower the logging level, the more tensor information you'll get. Feel free to play around!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from enum import IntEnum\n",
    "# Control how much debugging output we want\n",
    "class TensorLoggingLevels(IntEnum):\n",
    "    attention = 1\n",
    "    attention_head = 2\n",
    "    multihead_attention_block = 3\n",
    "    enc_dec_block = 4\n",
    "    enc_dec = 5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll be using an enum to refer to dimensions whenever possible to improve readability."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Dim(IntEnum):\n",
    "    batch = 0\n",
    "    seq = 1\n",
    "    feature = 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Components"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Scaled dot product attention"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Transformer is an attention-based architecture. The attention used in the Transformer is the scaled dot product attention, represented by the following formula."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$ \\textrm{Attention}(Q, K, V) = \\textrm{softmax}(\\frac{QK^T}{\\sqrt{d_k}})V $$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![image](https://i2.wp.com/mlexplained.com/wp-content/uploads/2017/12/scaled_dot_product_attention.png?zoom=2&w=750)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "\n",
    "class ScaledDotProductAttention(nn.Module):\n",
    "    level = TensorLoggingLevels.attention # Logging level: \n",
    "    def __init__(self, dropout=0.1):\n",
    "        super().__init__()\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, q, k, v, mask=None):\n",
    "        d_k = k.size(-1) # get the size of the key\n",
    "        assert q.size(-1) == d_k\n",
    "\n",
    "        # compute the dot product between queries and keys for\n",
    "        # each batch and position in the sequence\n",
    "        attn = torch.bmm(q, k.transpose(Dim.seq, Dim.feature)) # (Batch, Seq, Seq)\n",
    "        # we get an attention score between each position in the sequence\n",
    "        # for each batch\n",
    "\n",
    "        # scale the dot products by the dimensionality (see the paper for why we do this!)\n",
    "        attn = attn / math.sqrt(d_k)\n",
    "        # normalize the weights across the sequence dimension\n",
    "        # (Note that since we transposed, the sequence and feature dimensions are switched)\n",
    "        attn = torch.exp(attn)\n",
    "        log_size(attn, \"attention weight\") # (Batch, Seq, Seq)\n",
    "        \n",
    "        # fill attention weights with 0s where padded\n",
    "        if mask is not None: attn = attn.masked_fill(mask, 0)\n",
    "        attn = attn / attn.sum(dim=-1, keepdim=True)\n",
    "        attn = self.dropout(attn)\n",
    "        output = torch.bmm(attn, v) # (Batch, Seq, Feature)\n",
    "        log_size(output, \"attention output size\") # (Batch, Seq, Seq)\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "attn = ScaledDotProductAttention()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "q = torch.rand(5, 10, 20)\n",
    "k = torch.rand(5, 10, 20)\n",
    "v = torch.rand(5, 10, 20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[ScaledDotProductAttention] attention weight size=torch.Size([5, 10, 10])\n",
      "[ScaledDotProductAttention] attention output size size=torch.Size([5, 10, 20])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[[0.4483, 0.6249, 0.4408, 0.5156, 0.4743, 0.5172, 0.4808, 0.6804,\n",
       "          0.4991, 0.4802, 0.6905, 0.4754, 0.4845, 0.4595, 0.4427, 0.4852,\n",
       "          0.6439, 0.3720, 0.4924, 0.4870],\n",
       "         [0.4369, 0.5293, 0.4167, 0.4561, 0.4426, 0.4175, 0.4314, 0.6212,\n",
       "          0.4854, 0.3984, 0.5607, 0.3664, 0.4413, 0.3703, 0.3464, 0.4736,\n",
       "          0.5042, 0.3894, 0.3976, 0.4720],\n",
       "         [0.4000, 0.6714, 0.4945, 0.5342, 0.3825, 0.4558, 0.4701, 0.7120,\n",
       "          0.5406, 0.4414, 0.6396, 0.5229, 0.4415, 0.5028, 0.4552, 0.4466,\n",
       "          0.5253, 0.4450, 0.4930, 0.4459],\n",
       "         [0.4549, 0.6190, 0.4529, 0.5140, 0.4734, 0.5101, 0.4743, 0.6831,\n",
       "          0.4990, 0.4746, 0.6796, 0.4653, 0.4834, 0.4589, 0.4280, 0.4894,\n",
       "          0.6144, 0.3771, 0.4860, 0.4962],\n",
       "         [0.3270, 0.4722, 0.4131, 0.3552, 0.2582, 0.2866, 0.3154, 0.5365,\n",
       "          0.5085, 0.2973, 0.4904, 0.3656, 0.3859, 0.3883, 0.3579, 0.3699,\n",
       "          0.4034, 0.3194, 0.3842, 0.3813],\n",
       "         [0.5001, 0.7027, 0.5021, 0.5506, 0.4861, 0.5732, 0.5481, 0.7682,\n",
       "          0.5530, 0.5572, 0.7297, 0.5702, 0.5134, 0.5041, 0.5373, 0.5110,\n",
       "          0.6554, 0.4534, 0.5316, 0.4916],\n",
       "         [0.4578, 0.6273, 0.4483, 0.5206, 0.4632, 0.5092, 0.4768, 0.6956,\n",
       "          0.5191, 0.4748, 0.6856, 0.4704, 0.5005, 0.4617, 0.4378, 0.4991,\n",
       "          0.6385, 0.3871, 0.5010, 0.5012],\n",
       "         [0.4934, 0.6966, 0.5083, 0.5573, 0.4938, 0.5555, 0.5477, 0.7655,\n",
       "          0.5721, 0.5340, 0.7166, 0.5699, 0.5153, 0.5179, 0.5370, 0.5092,\n",
       "          0.6377, 0.4584, 0.5315, 0.4964],\n",
       "         [0.4469, 0.4997, 0.3586, 0.4489, 0.4678, 0.4884, 0.4380, 0.5993,\n",
       "          0.4375, 0.4227, 0.5741, 0.4251, 0.4901, 0.4417, 0.4130, 0.4258,\n",
       "          0.5537, 0.2982, 0.4623, 0.4523],\n",
       "         [0.4903, 0.6933, 0.5045, 0.5472, 0.5009, 0.5631, 0.5486, 0.7524,\n",
       "          0.5656, 0.5382, 0.7284, 0.5737, 0.5140, 0.5272, 0.5467, 0.4985,\n",
       "          0.6547, 0.4357, 0.5324, 0.4924]],\n",
       "\n",
       "        [[0.4519, 0.5578, 0.5542, 0.5214, 0.4943, 0.4578, 0.3808, 0.4553,\n",
       "          0.5869, 0.5293, 0.5423, 0.4184, 0.3226, 0.3980, 0.4553, 0.4766,\n",
       "          0.5001, 0.4874, 0.4138, 0.5982],\n",
       "         [0.5260, 0.6302, 0.6324, 0.5950, 0.5721, 0.5420, 0.4087, 0.5023,\n",
       "          0.6949, 0.6104, 0.6128, 0.4963, 0.4308, 0.5146, 0.5856, 0.5706,\n",
       "          0.6086, 0.6186, 0.5175, 0.6787],\n",
       "         [0.5314, 0.6171, 0.6321, 0.5972, 0.5772, 0.5449, 0.3960, 0.5018,\n",
       "          0.6969, 0.6076, 0.6212, 0.5030, 0.4314, 0.5131, 0.5725, 0.5608,\n",
       "          0.5988, 0.6161, 0.5087, 0.6798],\n",
       "         [0.4001, 0.5974, 0.5155, 0.4086, 0.4440, 0.4344, 0.4032, 0.3565,\n",
       "          0.5595, 0.4353, 0.4777, 0.3782, 0.3842, 0.3951, 0.4662, 0.4924,\n",
       "          0.5728, 0.5678, 0.4978, 0.4840],\n",
       "         [0.4588, 0.4575, 0.4764, 0.4745, 0.4986, 0.4666, 0.3127, 0.3126,\n",
       "          0.6443, 0.4742, 0.5617, 0.4360, 0.4160, 0.4057, 0.4493, 0.4412,\n",
       "          0.5163, 0.4737, 0.4321, 0.5188],\n",
       "         [0.5233, 0.6274, 0.6329, 0.5869, 0.5768, 0.5450, 0.4096, 0.4845,\n",
       "          0.7085, 0.5974, 0.6210, 0.5011, 0.4435, 0.5074, 0.5598, 0.5507,\n",
       "          0.6206, 0.6133, 0.5257, 0.6607],\n",
       "         [0.4410, 0.4548, 0.4872, 0.5306, 0.4410, 0.4848, 0.3529, 0.3664,\n",
       "          0.6244, 0.4558, 0.4749, 0.4771, 0.3335, 0.4231, 0.4261, 0.4352,\n",
       "          0.4266, 0.4856, 0.3597, 0.5045],\n",
       "         [0.4978, 0.5627, 0.5871, 0.5357, 0.5764, 0.5082, 0.3254, 0.4163,\n",
       "          0.6877, 0.5386, 0.6050, 0.4516, 0.4338, 0.4661, 0.4936, 0.4943,\n",
       "          0.6054, 0.5622, 0.5245, 0.6189],\n",
       "         [0.3755, 0.4908, 0.4193, 0.3715, 0.3792, 0.4018, 0.3851, 0.2618,\n",
       "          0.5333, 0.3786, 0.4408, 0.3675, 0.3597, 0.3522, 0.4378, 0.4491,\n",
       "          0.4952, 0.4857, 0.4204, 0.4089],\n",
       "         [0.5238, 0.6342, 0.6308, 0.6048, 0.5703, 0.5468, 0.4212, 0.5006,\n",
       "          0.7049, 0.6019, 0.6010, 0.5009, 0.4170, 0.5127, 0.5755, 0.5717,\n",
       "          0.6009, 0.6150, 0.5136, 0.6738]],\n",
       "\n",
       "        [[0.4895, 0.4341, 0.5660, 0.5691, 0.4895, 0.4109, 0.4156, 0.5709,\n",
       "          0.4254, 0.2942, 0.3869, 0.3793, 0.3086, 0.3911, 0.6697, 0.5181,\n",
       "          0.4237, 0.2983, 0.5160, 0.3884],\n",
       "         [0.5409, 0.5319, 0.7275, 0.5853, 0.5566, 0.5647, 0.4346, 0.6866,\n",
       "          0.4727, 0.4091, 0.4879, 0.4941, 0.3900, 0.5638, 0.6128, 0.5243,\n",
       "          0.4764, 0.3673, 0.5791, 0.5175],\n",
       "         [0.4889, 0.4557, 0.6307, 0.4842, 0.5428, 0.4258, 0.3615, 0.4847,\n",
       "          0.4058, 0.2944, 0.3876, 0.3449, 0.3482, 0.3895, 0.6296, 0.5052,\n",
       "          0.3816, 0.3534, 0.5722, 0.4619],\n",
       "         [0.5867, 0.5893, 0.7437, 0.6231, 0.6533, 0.5817, 0.5414, 0.7051,\n",
       "          0.5429, 0.4231, 0.4713, 0.4964, 0.4675, 0.5802, 0.7109, 0.6288,\n",
       "          0.5442, 0.3834, 0.6494, 0.5841],\n",
       "         [0.5752, 0.5732, 0.7302, 0.6255, 0.6460, 0.5820, 0.5607, 0.7048,\n",
       "          0.5622, 0.4353, 0.4591, 0.5052, 0.4765, 0.5924, 0.7106, 0.6403,\n",
       "          0.5505, 0.3692, 0.6500, 0.5883],\n",
       "         [0.5833, 0.5935, 0.7318, 0.6220, 0.6635, 0.5626, 0.5479, 0.7001,\n",
       "          0.5428, 0.4201, 0.4713, 0.4828, 0.4634, 0.5636, 0.7259, 0.6242,\n",
       "          0.5463, 0.3702, 0.6498, 0.5855],\n",
       "         [0.5830, 0.5938, 0.7495, 0.6212, 0.6546, 0.5782, 0.5376, 0.7143,\n",
       "          0.5398, 0.4257, 0.4946, 0.5021, 0.4505, 0.5791, 0.7144, 0.6148,\n",
       "          0.5488, 0.3843, 0.6672, 0.5988],\n",
       "         [0.5958, 0.5984, 0.7511, 0.6225, 0.6533, 0.5703, 0.5211, 0.7141,\n",
       "          0.5366, 0.4267, 0.4860, 0.4957, 0.4612, 0.5805, 0.7234, 0.6166,\n",
       "          0.5360, 0.3813, 0.6536, 0.5767],\n",
       "         [0.4682, 0.5139, 0.6436, 0.5328, 0.5830, 0.5258, 0.5158, 0.6636,\n",
       "          0.4730, 0.3958, 0.4547, 0.4622, 0.4137, 0.5484, 0.6125, 0.5391,\n",
       "          0.4954, 0.3062, 0.6018, 0.5511],\n",
       "         [0.5816, 0.5731, 0.7485, 0.6363, 0.6470, 0.5792, 0.5456, 0.7014,\n",
       "          0.5480, 0.4300, 0.4886, 0.4980, 0.4433, 0.5588, 0.7180, 0.6207,\n",
       "          0.5466, 0.3891, 0.6590, 0.6050]],\n",
       "\n",
       "        [[0.3381, 0.5657, 0.4135, 0.8694, 0.4851, 0.5170, 0.4996, 0.6773,\n",
       "          0.6416, 0.4715, 0.7882, 0.5475, 0.5995, 0.5806, 0.6181, 0.6671,\n",
       "          0.4540, 0.6698, 0.4332, 0.5073],\n",
       "         [0.3024, 0.5383, 0.3370, 0.7999, 0.4071, 0.5004, 0.4751, 0.5800,\n",
       "          0.6128, 0.4237, 0.7184, 0.4622, 0.5300, 0.5499, 0.5889, 0.6301,\n",
       "          0.4279, 0.6485, 0.3787, 0.4639],\n",
       "         [0.3206, 0.4735, 0.3720, 0.7676, 0.4732, 0.5030, 0.5008, 0.5787,\n",
       "          0.6731, 0.4568, 0.7020, 0.4753, 0.5733, 0.5784, 0.5607, 0.5973,\n",
       "          0.4390, 0.5662, 0.4275, 0.4675],\n",
       "         [0.3297, 0.5582, 0.4101, 0.8694, 0.4945, 0.5168, 0.4920, 0.6736,\n",
       "          0.6607, 0.4647, 0.8014, 0.5363, 0.5988, 0.5644, 0.6326, 0.6725,\n",
       "          0.4582, 0.6661, 0.4217, 0.5121],\n",
       "         [0.2700, 0.4502, 0.2906, 0.7233, 0.3760, 0.4946, 0.3559, 0.4475,\n",
       "          0.5303, 0.3559, 0.6429, 0.4358, 0.5250, 0.4668, 0.5306, 0.5428,\n",
       "          0.3369, 0.5286, 0.3471, 0.3474],\n",
       "         [0.3323, 0.5565, 0.4184, 0.8670, 0.4868, 0.5321, 0.4974, 0.6771,\n",
       "          0.6516, 0.4725, 0.7918, 0.5392, 0.6046, 0.5838, 0.6189, 0.6695,\n",
       "          0.4469, 0.6585, 0.4471, 0.5016],\n",
       "         [0.3414, 0.5684, 0.4098, 0.8682, 0.4744, 0.5025, 0.5073, 0.6762,\n",
       "          0.6483, 0.4763, 0.8012, 0.5386, 0.5958, 0.5779, 0.6276, 0.6738,\n",
       "          0.4402, 0.6543, 0.4419, 0.5115],\n",
       "         [0.3447, 0.5628, 0.4066, 0.8624, 0.4762, 0.5103, 0.4962, 0.6720,\n",
       "          0.6505, 0.4748, 0.7981, 0.5382, 0.6016, 0.5863, 0.6239, 0.6757,\n",
       "          0.4446, 0.6650, 0.4445, 0.5055],\n",
       "         [0.2334, 0.4831, 0.4002, 0.7897, 0.4963, 0.5093, 0.4063, 0.6217,\n",
       "          0.6214, 0.3969, 0.7214, 0.4925, 0.5375, 0.4781, 0.5722, 0.5860,\n",
       "          0.4236, 0.5882, 0.3693, 0.4570],\n",
       "         [0.3025, 0.4744, 0.3737, 0.7890, 0.4627, 0.4929, 0.3838, 0.5479,\n",
       "          0.5777, 0.4095, 0.7254, 0.5264, 0.5979, 0.4949, 0.5678, 0.5809,\n",
       "          0.3496, 0.5173, 0.4101, 0.3921]],\n",
       "\n",
       "        [[0.4132, 0.4699, 0.5051, 0.4588, 0.5631, 0.6553, 0.6388, 0.5417,\n",
       "          0.7043, 0.5008, 0.7782, 0.5163, 0.5241, 0.5968, 0.7099, 0.4928,\n",
       "          0.5146, 0.4108, 0.6716, 0.4133],\n",
       "         [0.4035, 0.4603, 0.5034, 0.4115, 0.4938, 0.6108, 0.5747, 0.4552,\n",
       "          0.6241, 0.4241, 0.7213, 0.4501, 0.4506, 0.5059, 0.6903, 0.4680,\n",
       "          0.5180, 0.3674, 0.6056, 0.3833],\n",
       "         [0.4282, 0.4625, 0.4958, 0.4542, 0.5557, 0.6629, 0.6351, 0.5468,\n",
       "          0.6909, 0.4881, 0.7705, 0.5126, 0.5277, 0.6063, 0.6983, 0.4750,\n",
       "          0.5221, 0.4025, 0.6783, 0.4044],\n",
       "         [0.3046, 0.3920, 0.4037, 0.3562, 0.4546, 0.5787, 0.4974, 0.3705,\n",
       "          0.5786, 0.3408, 0.5977, 0.4036, 0.4175, 0.4955, 0.5231, 0.4804,\n",
       "          0.3341, 0.3284, 0.5927, 0.3084],\n",
       "         [0.3985, 0.4252, 0.4889, 0.3610, 0.3970, 0.5182, 0.5069, 0.4049,\n",
       "          0.5592, 0.4249, 0.6257, 0.4759, 0.4692, 0.4685, 0.6301, 0.4556,\n",
       "          0.5164, 0.3884, 0.5600, 0.3441],\n",
       "         [0.3674, 0.4463, 0.3846, 0.4290, 0.5369, 0.5654, 0.5973, 0.4905,\n",
       "          0.6166, 0.3806, 0.6751, 0.4116, 0.4574, 0.4649, 0.5818, 0.4194,\n",
       "          0.4048, 0.3707, 0.6176, 0.3799],\n",
       "         [0.4130, 0.4518, 0.4789, 0.3746, 0.4768, 0.5784, 0.5584, 0.4471,\n",
       "          0.6237, 0.4701, 0.6683, 0.4863, 0.5115, 0.4734, 0.6156, 0.4929,\n",
       "          0.4978, 0.3824, 0.5797, 0.3892],\n",
       "         [0.4316, 0.4865, 0.5152, 0.4473, 0.5558, 0.6637, 0.6290, 0.5232,\n",
       "          0.6984, 0.4788, 0.7552, 0.5168, 0.5380, 0.5778, 0.6981, 0.5187,\n",
       "          0.5242, 0.4109, 0.6752, 0.4030],\n",
       "         [0.4100, 0.4439, 0.4819, 0.3704, 0.4641, 0.5662, 0.5372, 0.4295,\n",
       "          0.6113, 0.4673, 0.6542, 0.4966, 0.5094, 0.4646, 0.5977, 0.5005,\n",
       "          0.4877, 0.3886, 0.5684, 0.3868],\n",
       "         [0.3548, 0.3753, 0.3208, 0.3576, 0.4488, 0.5363, 0.5481, 0.4679,\n",
       "          0.5282, 0.3175, 0.5969, 0.3331, 0.4097, 0.4760, 0.5394, 0.3062,\n",
       "          0.4079, 0.2874, 0.5813, 0.2970]]])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "attn(q, k, v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Multi-Head Attention"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we turn to the core component in the Transformer architecture: the multi-head attention block. This block applies linear transformations to the input, then applies scaled dot product attention."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![image](https://i2.wp.com/mlexplained.com/wp-content/uploads/2017/12/multi_head_attention.png?zoom=2&resize=224%2C293)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AttentionHead(nn.Module):\n",
    "    level = TensorLoggingLevels.attention_head\n",
    "    def __init__(self, d_model, d_feature, dropout=0.1):\n",
    "        super().__init__()\n",
    "        # We will assume the queries, keys, and values all have the same feature size\n",
    "        self.attn = ScaledDotProductAttention(dropout)\n",
    "        self.query_tfm = nn.Linear(d_model, d_feature)\n",
    "        self.key_tfm = nn.Linear(d_model, d_feature)\n",
    "        self.value_tfm = nn.Linear(d_model, d_feature)\n",
    "\n",
    "    def forward(self, queries, keys, values, mask=None):\n",
    "        Q = self.query_tfm(queries) # (Batch, Seq, Feature)\n",
    "        K = self.key_tfm(keys) # (Batch, Seq, Feature)\n",
    "        V = self.value_tfm(values) # (Batch, Seq, Feature)\n",
    "        log_size(Q, \"queries, keys, vals\")\n",
    "        # compute multiple attention weighted sums\n",
    "        x = self.attn(Q, K, V)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[AttentionHead] queries, keys, vals size=torch.Size([5, 10, 20])\n",
      "[ScaledDotProductAttention] attention weight size=torch.Size([5, 10, 10])\n",
      "[ScaledDotProductAttention] attention output size size=torch.Size([5, 10, 20])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[[ 0.0238,  0.1517, -0.1136,  0.3600,  0.5793,  0.4810,  0.5199,\n",
       "           0.2519,  0.4256,  0.4177, -0.2117, -0.6132,  0.0139, -0.1483,\n",
       "           0.0267,  0.1633, -0.0067,  0.8864, -0.3889, -0.5447],\n",
       "         [ 0.0454,  0.1453, -0.0759,  0.3506,  0.5473,  0.4195,  0.4682,\n",
       "           0.2472,  0.4111,  0.3637, -0.2134, -0.5537,  0.0039, -0.1573,\n",
       "           0.0459,  0.1266,  0.0056,  0.8011, -0.3324, -0.4825],\n",
       "         [ 0.0131,  0.1344, -0.0980,  0.3264,  0.5227,  0.4492,  0.4796,\n",
       "           0.2134,  0.3958,  0.3693, -0.1795, -0.5554,  0.0123, -0.1559,\n",
       "           0.0100,  0.1447, -0.0218,  0.7680, -0.3628, -0.4881],\n",
       "         [-0.0106,  0.1171, -0.1473,  0.3370,  0.5457,  0.5062,  0.4814,\n",
       "           0.2054,  0.3945,  0.3926, -0.1718, -0.5806,  0.0201, -0.1544,\n",
       "           0.0344,  0.1406, -0.0233,  0.8585, -0.4306, -0.5781],\n",
       "         [-0.0053,  0.1495, -0.1153,  0.3350,  0.5565,  0.5025,  0.5206,\n",
       "           0.1967,  0.4121,  0.4178, -0.1702, -0.6003,  0.0268, -0.1763,\n",
       "           0.0126,  0.1633, -0.0210,  0.8780, -0.4316, -0.5706],\n",
       "         [ 0.0072,  0.1682, -0.1235,  0.3711,  0.6119,  0.5301,  0.5604,\n",
       "           0.2358,  0.4413,  0.4615, -0.2037, -0.6570,  0.0260, -0.1702,\n",
       "           0.0277,  0.1803, -0.0056,  0.9931, -0.4541, -0.6247],\n",
       "         [-0.0019,  0.1042, -0.1523,  0.2775,  0.4832,  0.4330,  0.4114,\n",
       "           0.1905,  0.3542,  0.3672, -0.1569, -0.5036,  0.0275, -0.1246,\n",
       "           0.0564,  0.1260, -0.0058,  0.7603, -0.3635, -0.4972],\n",
       "         [ 0.0068,  0.1665, -0.1256,  0.3716,  0.6126,  0.5312,  0.5581,\n",
       "           0.2365,  0.4423,  0.4616, -0.2034, -0.6569,  0.0252, -0.1707,\n",
       "           0.0303,  0.1784, -0.0038,  0.9938, -0.4544, -0.6257],\n",
       "         [-0.0023,  0.1271, -0.1078,  0.2707,  0.4551,  0.4136,  0.4076,\n",
       "           0.1648,  0.3295,  0.3635, -0.1582, -0.4913, -0.0103, -0.1363,\n",
       "          -0.0055,  0.1553,  0.0014,  0.6735, -0.3006, -0.4176],\n",
       "         [ 0.0067,  0.1655, -0.1264,  0.3708,  0.6114,  0.5312,  0.5579,\n",
       "           0.2356,  0.4405,  0.4598, -0.2017, -0.6558,  0.0257, -0.1709,\n",
       "           0.0290,  0.1785, -0.0039,  0.9929, -0.4561, -0.6249]],\n",
       "\n",
       "        [[ 0.0196,  0.2096, -0.1021,  0.2403,  0.5438,  0.4413,  0.4342,\n",
       "           0.2121,  0.3255,  0.3758, -0.1650, -0.5612,  0.0443, -0.1916,\n",
       "           0.0663,  0.2187,  0.0112,  0.8601, -0.3491, -0.4969],\n",
       "         [ 0.0338,  0.1914, -0.0799,  0.1776,  0.5023,  0.3993,  0.3809,\n",
       "           0.2147,  0.2945,  0.3772, -0.1554, -0.4834,  0.0176, -0.1902,\n",
       "           0.0230,  0.1528,  0.0541,  0.8283, -0.2926, -0.4276],\n",
       "         [ 0.0099,  0.2011, -0.1319,  0.2427,  0.5574,  0.4854,  0.4572,\n",
       "           0.2062,  0.3207,  0.3631, -0.1631, -0.5296,  0.0417, -0.1677,\n",
       "           0.0557,  0.1665, -0.0049,  0.8964, -0.4099, -0.4968],\n",
       "         [ 0.0101,  0.1874, -0.1206,  0.2186,  0.5413,  0.4473,  0.4182,\n",
       "           0.2085,  0.3329,  0.3885, -0.1656, -0.5448,  0.0281, -0.1608,\n",
       "           0.0607,  0.1897,  0.0274,  0.8362, -0.3554, -0.4961],\n",
       "         [ 0.0138,  0.2223, -0.1301,  0.2540,  0.6194,  0.5215,  0.4744,\n",
       "           0.2334,  0.3586,  0.4320, -0.1840, -0.6031,  0.0312, -0.2043,\n",
       "           0.0504,  0.2142,  0.0302,  0.9817, -0.4051, -0.5527],\n",
       "         [-0.0073,  0.1797, -0.1531,  0.1999,  0.4950,  0.4575,  0.3928,\n",
       "           0.1557,  0.2730,  0.3228, -0.1443, -0.4458,  0.0253, -0.1414,\n",
       "           0.0123,  0.1586, -0.0247,  0.7914, -0.3645, -0.4629],\n",
       "         [ 0.0288,  0.1798, -0.0905,  0.1890,  0.4448,  0.3751,  0.3776,\n",
       "           0.1828,  0.2600,  0.2764, -0.1257, -0.4245,  0.0431, -0.1361,\n",
       "           0.0429,  0.1482, -0.0082,  0.7114, -0.3051, -0.3560],\n",
       "         [ 0.0143,  0.2221, -0.1298,  0.2510,  0.6200,  0.5206,  0.4734,\n",
       "           0.2332,  0.3590,  0.4332, -0.1839, -0.6007,  0.0315, -0.2043,\n",
       "           0.0489,  0.2128,  0.0307,  0.9804, -0.4038, -0.5499],\n",
       "         [ 0.0190,  0.2104, -0.1025,  0.2432,  0.5451,  0.4443,  0.4353,\n",
       "           0.2131,  0.3259,  0.3767, -0.1655, -0.5655,  0.0429, -0.1928,\n",
       "           0.0671,  0.2209,  0.0122,  0.8648, -0.3507, -0.5013],\n",
       "         [ 0.0144,  0.2222, -0.1283,  0.2535,  0.6202,  0.5219,  0.4732,\n",
       "           0.2345,  0.3586,  0.4327, -0.1841, -0.6043,  0.0293, -0.2049,\n",
       "           0.0509,  0.2133,  0.0341,  0.9829, -0.4051, -0.5513]],\n",
       "\n",
       "        [[ 0.0295,  0.1740, -0.0759,  0.2066,  0.4570,  0.3266,  0.3229,\n",
       "           0.1910,  0.3346,  0.2877, -0.1449, -0.4485,  0.0378, -0.1347,\n",
       "           0.0390,  0.2260, -0.0099,  0.6830, -0.2188, -0.4221],\n",
       "         [ 0.0695,  0.1963, -0.0590,  0.2935,  0.5360,  0.3932,  0.4835,\n",
       "           0.2209,  0.4094,  0.3183, -0.1732, -0.6056,  0.0686, -0.1676,\n",
       "           0.0472,  0.2661, -0.0312,  0.8378, -0.3650, -0.5253],\n",
       "         [ 0.0626,  0.2355, -0.0942,  0.3293,  0.6414,  0.4640,  0.5387,\n",
       "           0.2516,  0.4655,  0.3740, -0.2144, -0.6486,  0.0715, -0.2154,\n",
       "           0.0429,  0.2821, -0.0539,  0.9666, -0.4009, -0.5858],\n",
       "         [ 0.0655,  0.2266, -0.0647,  0.2857,  0.5895,  0.3987,  0.4908,\n",
       "           0.2249,  0.3906,  0.3091, -0.1982, -0.5653,  0.0704, -0.2167,\n",
       "           0.0483,  0.2413, -0.0349,  0.8769, -0.3815, -0.4995],\n",
       "         [ 0.0550,  0.2268, -0.1113,  0.3222,  0.6000,  0.4371,  0.4933,\n",
       "           0.2303,  0.4013,  0.3219, -0.2165, -0.5701,  0.0588, -0.1674,\n",
       "           0.0439,  0.2716, -0.0603,  0.8603, -0.3569, -0.5228],\n",
       "         [ 0.0618,  0.2346, -0.0924,  0.3282,  0.6401,  0.4628,  0.5382,\n",
       "           0.2510,  0.4656,  0.3750, -0.2116, -0.6484,  0.0725, -0.2188,\n",
       "           0.0407,  0.2807, -0.0538,  0.9694, -0.4008, -0.5858],\n",
       "         [ 0.0638,  0.2344, -0.0905,  0.3252,  0.6372,  0.4598,  0.5350,\n",
       "           0.2499,  0.4675,  0.3767, -0.2106, -0.6502,  0.0728, -0.2170,\n",
       "           0.0412,  0.2837, -0.0528,  0.9640, -0.3958, -0.5834],\n",
       "         [ 0.0499,  0.2144, -0.0666,  0.2904,  0.5751,  0.4004,  0.4585,\n",
       "           0.2228,  0.4261,  0.3627, -0.1824, -0.5890,  0.0643, -0.2185,\n",
       "           0.0422,  0.2583, -0.0331,  0.8811, -0.3258, -0.5211],\n",
       "         [ 0.0542,  0.2250, -0.1112,  0.3242,  0.5974,  0.4361,  0.4934,\n",
       "           0.2304,  0.3994,  0.3198, -0.2159, -0.5693,  0.0585, -0.1649,\n",
       "           0.0445,  0.2710, -0.0598,  0.8578, -0.3573, -0.5233],\n",
       "         [ 0.0696,  0.1924, -0.0565,  0.2872,  0.5272,  0.3871,  0.4780,\n",
       "           0.2170,  0.4078,  0.3172, -0.1680, -0.6000,  0.0697, -0.1687,\n",
       "           0.0430,  0.2622, -0.0327,  0.8286, -0.3601, -0.5189]],\n",
       "\n",
       "        [[ 0.1008,  0.1332, -0.1045,  0.2755,  0.4868,  0.3884,  0.5732,\n",
       "           0.2722,  0.3364,  0.3115, -0.1656, -0.4898,  0.1101, -0.0723,\n",
       "          -0.0025,  0.1130, -0.0098,  0.8464, -0.4462, -0.4206],\n",
       "         [ 0.0888,  0.1158, -0.1141,  0.3599,  0.5650,  0.4601,  0.6857,\n",
       "           0.3534,  0.4279,  0.4090, -0.1733, -0.6467,  0.1418, -0.0726,\n",
       "          -0.0078,  0.1146, -0.0069,  1.0460, -0.5350, -0.5701],\n",
       "         [ 0.0878,  0.0817, -0.1089,  0.2783,  0.4647,  0.3947,  0.5224,\n",
       "           0.3098,  0.3430,  0.3327, -0.1241, -0.5164,  0.1014, -0.0514,\n",
       "          -0.0110,  0.0969,  0.0387,  0.8541, -0.3977, -0.4561],\n",
       "         [ 0.0549,  0.0805, -0.0862,  0.2939,  0.4326,  0.3386,  0.5632,\n",
       "           0.2860,  0.3434,  0.3472, -0.1440, -0.5333,  0.1173, -0.0616,\n",
       "          -0.0137,  0.0916, -0.0113,  0.8477, -0.4285, -0.4836],\n",
       "         [ 0.0813,  0.0841, -0.1359,  0.3302,  0.5053,  0.4369,  0.6061,\n",
       "           0.3017,  0.3765,  0.3566, -0.1460, -0.5657,  0.1313, -0.0754,\n",
       "          -0.0137,  0.0955, -0.0207,  0.9095, -0.4838, -0.5146],\n",
       "         [ 0.0869,  0.1143, -0.1151,  0.3636,  0.5669,  0.4624,  0.6836,\n",
       "           0.3546,  0.4288,  0.4106, -0.1734, -0.6486,  0.1411, -0.0723,\n",
       "          -0.0057,  0.1146, -0.0057,  1.0460, -0.5332, -0.5719],\n",
       "         [ 0.0852,  0.1118, -0.1142,  0.3632,  0.5663,  0.4621,  0.6811,\n",
       "           0.3576,  0.4305,  0.4106, -0.1692, -0.6506,  0.1416, -0.0702,\n",
       "          -0.0053,  0.1155, -0.0014,  1.0475, -0.5316, -0.5738],\n",
       "         [ 0.0804,  0.0835, -0.1356,  0.3287,  0.5026,  0.4350,  0.6040,\n",
       "           0.3006,  0.3745,  0.3560, -0.1448, -0.5640,  0.1313, -0.0744,\n",
       "          -0.0139,  0.0959, -0.0207,  0.9067, -0.4820, -0.5132],\n",
       "         [ 0.1147,  0.1317, -0.0748,  0.2280,  0.4174,  0.3150,  0.5787,\n",
       "           0.2717,  0.3206,  0.3443, -0.1406, -0.5094,  0.1225, -0.0557,\n",
       "          -0.0338,  0.1225,  0.0151,  0.8479, -0.4201, -0.4106],\n",
       "         [ 0.0609,  0.0706, -0.1220,  0.3281,  0.4886,  0.4082,  0.6062,\n",
       "           0.3304,  0.3989,  0.3749, -0.1265, -0.6053,  0.1371, -0.0431,\n",
       "          -0.0060,  0.1301,  0.0112,  0.9317, -0.4744, -0.5501]],\n",
       "\n",
       "        [[ 0.0650,  0.1742, -0.1863,  0.2665,  0.6489,  0.5514,  0.5600,\n",
       "           0.2656,  0.4411,  0.4030, -0.1563, -0.5507,  0.0703, -0.1360,\n",
       "          -0.0155,  0.2206,  0.0552,  1.0062, -0.4609, -0.6206],\n",
       "         [ 0.0676,  0.1707, -0.1894,  0.2659,  0.6478,  0.5531,  0.5620,\n",
       "           0.2670,  0.4422,  0.3975, -0.1531, -0.5488,  0.0720, -0.1354,\n",
       "          -0.0183,  0.2192,  0.0537,  1.0031, -0.4637, -0.6209],\n",
       "         [ 0.0857,  0.1600, -0.1729,  0.2222,  0.5902,  0.4866,  0.5247,\n",
       "           0.2510,  0.4197,  0.3594, -0.1531, -0.4745,  0.0749, -0.1264,\n",
       "          -0.0215,  0.1774,  0.0295,  0.8992, -0.4078, -0.5510],\n",
       "         [ 0.0390,  0.1536, -0.1249,  0.1593,  0.4456,  0.3803,  0.3885,\n",
       "           0.1438,  0.2852,  0.3186, -0.1018, -0.3919,  0.0623, -0.1142,\n",
       "          -0.0206,  0.2217,  0.0457,  0.7221, -0.3171, -0.4238],\n",
       "         [ 0.0388,  0.1530, -0.1558,  0.2340,  0.5651,  0.4856,  0.4645,\n",
       "           0.2181,  0.3510,  0.3558, -0.1389, -0.4853,  0.0418, -0.1276,\n",
       "          -0.0209,  0.1935,  0.0651,  0.8894, -0.4016, -0.5536],\n",
       "         [ 0.0353,  0.1137, -0.1578,  0.2558,  0.5247,  0.4621,  0.4978,\n",
       "           0.2646,  0.3581,  0.2831, -0.0992, -0.4641,  0.0785, -0.0744,\n",
       "          -0.0030,  0.1536,  0.0420,  0.8404, -0.4312, -0.5275],\n",
       "         [ 0.0342,  0.1302, -0.1818,  0.2342,  0.5695,  0.4934,  0.5048,\n",
       "           0.2581,  0.3886,  0.3277, -0.1010, -0.4850,  0.0879, -0.0887,\n",
       "          -0.0324,  0.1971,  0.0374,  0.9015, -0.4368, -0.5673],\n",
       "         [ 0.0447,  0.1794, -0.1633,  0.2314,  0.5892,  0.4851,  0.4911,\n",
       "           0.2276,  0.3840,  0.3789, -0.1655, -0.4929,  0.0436, -0.1340,\n",
       "          -0.0065,  0.1998,  0.0528,  0.9032, -0.3988, -0.5625],\n",
       "         [ 0.0628,  0.1738, -0.1779,  0.2139,  0.5635,  0.4934,  0.4970,\n",
       "           0.2022,  0.3771,  0.3629, -0.1388, -0.4970,  0.0520, -0.1414,\n",
       "          -0.0402,  0.2334,  0.0420,  0.8810, -0.4103, -0.5553],\n",
       "         [ 0.0653,  0.1733, -0.1858,  0.2683,  0.6495,  0.5524,  0.5611,\n",
       "           0.2671,  0.4417,  0.4036, -0.1555, -0.5518,  0.0715, -0.1350,\n",
       "          -0.0135,  0.2208,  0.0570,  1.0081, -0.4619, -0.6209]]],\n",
       "       grad_fn=<BmmBackward>)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "attn_head = AttentionHead(20, 20)\n",
    "attn_head(q, k, v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The multi-head attention block simply applies multiple attention heads, then concatenates the outputs and applies a single linear projection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We'll supress logging from the scaled dot product attention now\n",
    "logger.setLevel(TensorLoggingLevels.attention_head)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadAttention(nn.Module):\n",
    "    level = TensorLoggingLevels.multihead_attention_block\n",
    "    def __init__(self, d_model, d_feature, n_heads, dropout=0.1):\n",
    "        super().__init__()\n",
    "        self.d_model = d_model\n",
    "        self.d_feature = d_feature\n",
    "        self.n_heads = n_heads\n",
    "        # in practice, d_model == d_feature * n_heads\n",
    "        assert d_model == d_feature * n_heads\n",
    "\n",
    "        # Note that this is very inefficient:\n",
    "        # I am merely implementing the heads separately because it is \n",
    "        # easier to understand this way\n",
    "        self.attn_heads = nn.ModuleList([\n",
    "            AttentionHead(d_model, d_feature, dropout) for _ in range(n_heads)\n",
    "        ])\n",
    "        self.projection = nn.Linear(d_feature * n_heads, d_model) \n",
    "    \n",
    "    def forward(self, queries, keys, values, mask=None):\n",
    "        log_size(queries, \"Input queries\")\n",
    "        x = [attn(queries, keys, values, mask=mask) # (Batch, Seq, Feature)\n",
    "             for i, attn in enumerate(self.attn_heads)]\n",
    "        log_size(x[0], \"output of single head\")\n",
    "        \n",
    "        # reconcatenate\n",
    "        x = torch.cat(x, dim=Dim.feature) # (Batch, Seq, D_Feature * n_heads)\n",
    "        log_size(x, \"concatenated output\")\n",
    "        x = self.projection(x) # (Batch, Seq, D_Model)\n",
    "        log_size(x, \"projected output\")\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[MultiHeadAttention] Input queries size=torch.Size([5, 10, 160])\n",
      "[AttentionHead] queries, keys, vals size=torch.Size([5, 10, 20])\n",
      "[AttentionHead] queries, keys, vals size=torch.Size([5, 10, 20])\n",
      "[AttentionHead] queries, keys, vals size=torch.Size([5, 10, 20])\n",
      "[AttentionHead] queries, keys, vals size=torch.Size([5, 10, 20])\n",
      "[AttentionHead] queries, keys, vals size=torch.Size([5, 10, 20])\n",
      "[AttentionHead] queries, keys, vals size=torch.Size([5, 10, 20])\n",
      "[AttentionHead] queries, keys, vals size=torch.Size([5, 10, 20])\n",
      "[AttentionHead] queries, keys, vals size=torch.Size([5, 10, 20])\n",
      "[MultiHeadAttention] output of single head size=torch.Size([5, 10, 20])\n",
      "[MultiHeadAttention] concatenated output size=torch.Size([5, 10, 160])\n",
      "[MultiHeadAttention] projected output size=torch.Size([5, 10, 160])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[[-5.5434e-02,  7.3836e-02,  2.1540e-01,  ...,  9.5668e-02,\n",
       "           2.5618e-01,  3.3740e-02],\n",
       "         [-5.1989e-02,  9.7644e-02,  1.7390e-01,  ...,  1.2145e-01,\n",
       "           2.1276e-01,  5.6385e-02],\n",
       "         [-5.7726e-02,  7.8944e-02,  2.1749e-01,  ...,  8.0804e-02,\n",
       "           2.5118e-01,  1.2463e-02],\n",
       "         ...,\n",
       "         [-6.4274e-02,  5.6929e-02,  1.9471e-01,  ...,  1.0194e-01,\n",
       "           2.4769e-01,  4.6319e-02],\n",
       "         [-3.4287e-02,  8.8949e-02,  2.4115e-01,  ...,  1.0646e-01,\n",
       "           1.9937e-01,  3.4541e-02],\n",
       "         [-5.6988e-02,  9.4813e-02,  2.2184e-01,  ...,  1.0748e-01,\n",
       "           2.5925e-01,  6.2197e-03]],\n",
       "\n",
       "        [[-1.7362e-03,  1.3871e-01,  2.2642e-01,  ...,  1.0594e-01,\n",
       "           2.1807e-01,  3.7778e-02],\n",
       "         [-1.9615e-02,  1.2238e-01,  2.5283e-01,  ...,  1.1735e-01,\n",
       "           1.8682e-01,  4.4808e-02],\n",
       "         [-1.8819e-02,  1.2427e-01,  2.6472e-01,  ...,  1.2316e-01,\n",
       "           2.0093e-01,  5.4456e-02],\n",
       "         ...,\n",
       "         [-4.7275e-02,  1.3365e-01,  2.5587e-01,  ...,  1.3074e-01,\n",
       "           1.7753e-01,  1.5566e-02],\n",
       "         [-2.6786e-02,  1.3173e-01,  2.5589e-01,  ...,  1.2248e-01,\n",
       "           2.0869e-01,  2.3813e-02],\n",
       "         [-9.0349e-02,  1.2743e-01,  1.7026e-01,  ...,  1.2832e-01,\n",
       "           2.3952e-01,  5.6115e-02]],\n",
       "\n",
       "        [[-3.2307e-02,  4.4406e-02,  2.0854e-01,  ...,  1.4595e-01,\n",
       "           2.1808e-01,  3.8419e-04],\n",
       "         [-4.2970e-02,  1.2763e-02,  2.0558e-01,  ...,  1.4893e-01,\n",
       "           2.2219e-01, -2.2180e-02],\n",
       "         [-6.5011e-02,  2.4007e-02,  2.1337e-01,  ...,  1.8114e-01,\n",
       "           2.3151e-01, -3.4369e-02],\n",
       "         ...,\n",
       "         [-4.1502e-02,  4.8484e-02,  1.8406e-01,  ...,  1.5471e-01,\n",
       "           2.0851e-01, -2.3843e-02],\n",
       "         [-4.8125e-02,  4.9388e-02,  1.9939e-01,  ...,  1.3571e-01,\n",
       "           2.6702e-01, -2.1225e-02],\n",
       "         [-5.4713e-02,  4.6316e-02,  2.2017e-01,  ...,  1.2882e-01,\n",
       "           2.4194e-01, -1.3253e-03]],\n",
       "\n",
       "        [[-9.7700e-02,  9.4153e-02,  2.8414e-01,  ...,  1.1452e-01,\n",
       "           2.8944e-01,  2.7665e-02],\n",
       "         [-7.1545e-02,  1.3726e-01,  3.4052e-01,  ...,  1.1537e-01,\n",
       "           2.0142e-01,  4.9502e-02],\n",
       "         [-1.1585e-01,  1.1944e-01,  2.7515e-01,  ...,  1.4937e-01,\n",
       "           2.1817e-01,  4.7201e-02],\n",
       "         ...,\n",
       "         [-9.1791e-02,  1.1561e-01,  2.8525e-01,  ...,  1.0322e-01,\n",
       "           2.6468e-01,  6.1614e-02],\n",
       "         [-8.0859e-02,  1.1107e-01,  3.0561e-01,  ...,  9.9594e-02,\n",
       "           2.4221e-01,  4.1798e-02],\n",
       "         [-7.2943e-02,  1.2736e-01,  3.3810e-01,  ...,  9.3637e-02,\n",
       "           2.3529e-01,  5.4366e-02]],\n",
       "\n",
       "        [[-6.0273e-02,  1.0379e-01,  2.2802e-01,  ...,  1.4011e-01,\n",
       "           2.2384e-01, -3.0530e-03],\n",
       "         [-3.2069e-02,  1.1405e-01,  2.4305e-01,  ...,  1.2542e-01,\n",
       "           2.3484e-01,  6.8733e-03],\n",
       "         [-4.1344e-02,  1.1123e-01,  2.2907e-01,  ...,  1.3332e-01,\n",
       "           2.2788e-01, -1.4862e-02],\n",
       "         ...,\n",
       "         [-5.2076e-02,  1.1014e-01,  2.2164e-01,  ...,  1.4577e-01,\n",
       "           2.2790e-01, -4.5416e-02],\n",
       "         [-3.1141e-02,  1.4255e-01,  2.3478e-01,  ...,  1.1450e-01,\n",
       "           2.4965e-01, -2.2048e-02],\n",
       "         [-4.2248e-02,  1.4350e-01,  2.7829e-01,  ...,  1.4013e-01,\n",
       "           2.1632e-01,  1.2843e-04]]], grad_fn=<AddBackward0>)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "heads = MultiHeadAttention(20 * 8, 20, 8)\n",
    "heads(q.repeat(1, 1, 8), \n",
    "      k.repeat(1, 1, 8), \n",
    "      v.repeat(1, 1, 8))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The Encoder"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With these core components in place, implementing the encoder is pretty easy."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![image](https://i2.wp.com/mlexplained.com/wp-content/uploads/2017/12/%E3%82%B9%E3%82%AF%E3%83%AA%E3%83%BC%E3%83%B3%E3%82%B7%E3%83%A7%E3%83%83%E3%83%88-2017-12-29-19.14.41.png?w=273)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The encoder consists of the following components:\n",
    "- A multi-head attention block\n",
    "- A simple feedforward neural network\n",
    "\n",
    "These components are connected using residual connections and layer normalization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We'll supress logging from the individual attention heads\n",
    "logger.setLevel(TensorLoggingLevels.multihead_attention_block)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Layer normalization is similar to batch normalization, but normalizes across the feature dimension instead of the batch dimension."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![image](https://i1.wp.com/mlexplained.com/wp-content/uploads/2018/01/%E3%82%B9%E3%82%AF%E3%83%AA%E3%83%BC%E3%83%B3%E3%82%B7%E3%83%A7%E3%83%83%E3%83%88-2018-01-11-11.48.12.png?w=1500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LayerNorm(nn.Module):\n",
    "    def __init__(self, d_model, eps=1e-8):\n",
    "        super(LayerNorm, self).__init__()\n",
    "        self.gamma = nn.Parameter(torch.ones(d_model))\n",
    "        self.beta = nn.Parameter(torch.zeros(d_model))\n",
    "        self.eps = eps\n",
    "\n",
    "    def forward(self, x):\n",
    "        mean = x.mean(-1, keepdim=True)\n",
    "        std = x.std(-1, keepdim=True)\n",
    "        return self.gamma * (x - mean) / (std + self.eps) + self.beta"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The encoder just stacks these together"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EncoderBlock(nn.Module):\n",
    "    level = TensorLoggingLevels.enc_dec_block\n",
    "    def __init__(self, d_model=512, d_feature=64,\n",
    "                 d_ff=2048, n_heads=8, dropout=0.1):\n",
    "        super().__init__()\n",
    "        self.attn_head = MultiHeadAttention(d_model, d_feature, n_heads, dropout)\n",
    "        self.layer_norm1 = LayerNorm(d_model)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        self.position_wise_feed_forward = nn.Sequential(\n",
    "            nn.Linear(d_model, d_ff),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(d_ff, d_model),\n",
    "        )\n",
    "        self.layer_norm2 = LayerNorm(d_model)\n",
    "        \n",
    "    def forward(self, x, mask=None):\n",
    "        log_size(x, \"Encoder block input\")\n",
    "        att = self.attn_head(x, x, x, mask=mask)\n",
    "        log_size(x, \"Attention output\")\n",
    "        # Apply normalization and residual connection\n",
    "        x = x + self.dropout(self.layer_norm1(att))\n",
    "        # Apply position-wise feedforward network\n",
    "        pos = self.position_wise_feed_forward(x)\n",
    "        log_size(x, \"Feedforward output\")\n",
    "        # Apply normalization and residual connection\n",
    "        x = x + self.dropout(self.layer_norm2(pos))\n",
    "        log_size(x, \"Encoder size output\")\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "enc = EncoderBlock()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[EncoderBlock] Encoder block input size=torch.Size([5, 10, 512])\n",
      "[MultiHeadAttention] Input queries size=torch.Size([5, 10, 512])\n",
      "[MultiHeadAttention] output of single head size=torch.Size([5, 10, 64])\n",
      "[MultiHeadAttention] concatenated output size=torch.Size([5, 10, 512])\n",
      "[MultiHeadAttention] projected output size=torch.Size([5, 10, 512])\n",
      "[EncoderBlock] Attention output size=torch.Size([5, 10, 512])\n",
      "[EncoderBlock] Feedforward output size=torch.Size([5, 10, 512])\n",
      "[EncoderBlock] Encoder size output size=torch.Size([5, 10, 512])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[[-0.9502, -0.0386,  2.0490,  ...,  4.3722, -0.8155, -0.6273],\n",
       "         [-1.1045, -0.5730,  1.6829,  ...,  3.7669, -0.6746,  0.1586],\n",
       "         [-0.5030, -0.8161,  1.9470,  ...,  3.7542, -1.6728, -0.6087],\n",
       "         ...,\n",
       "         [-1.6413, -0.8944,  1.5399,  ...,  5.0267, -0.0814,  0.7338],\n",
       "         [-1.1089, -0.3992,  1.6307,  ...,  3.5277, -0.9177,  0.6617],\n",
       "         [-0.7696, -0.6927,  0.7905,  ...,  3.7786, -1.1211, -0.0751]],\n",
       "\n",
       "        [[-0.3022,  0.1411,  2.5665,  ...,  4.0616, -0.1462,  0.6828],\n",
       "         [-0.8449,  0.0744,  1.3892,  ...,  4.0403, -0.6635,  0.3824],\n",
       "         [-0.8907,  0.0267,  2.7662,  ...,  3.9748, -0.4940,  0.9116],\n",
       "         ...,\n",
       "         [-0.1639,  0.4431,  0.7396,  ...,  3.4216, -1.1346,  0.2292],\n",
       "         [-1.2960,  0.0809,  2.0951,  ...,  4.3827, -1.1978,  0.9659],\n",
       "         [-0.8447, -0.1919,  2.9091,  ...,  4.0368, -1.3761,  1.0358]],\n",
       "\n",
       "        [[-2.2698, -0.7820,  1.9180,  ...,  3.3510, -0.1515,  0.4817],\n",
       "         [-1.7525, -0.6157,  1.3984,  ...,  4.2632, -0.6307,  0.3254],\n",
       "         [-1.2098,  0.6220,  2.1032,  ...,  4.1645, -0.8482,  0.0079],\n",
       "         ...,\n",
       "         [-2.2998, -0.0431,  1.4054,  ...,  3.8315, -1.6290, -0.1561],\n",
       "         [-1.4979, -0.2183,  2.4717,  ...,  3.2129, -0.5803, -0.1675],\n",
       "         [-1.4784, -0.5552,  2.1578,  ...,  3.3279, -1.1288, -0.1853]],\n",
       "\n",
       "        [[-0.9930,  0.7086,  2.2540,  ...,  4.2661, -0.8127,  0.7362],\n",
       "         [-1.0314, -0.5218,  1.7270,  ...,  3.5672,  0.5031,  0.1702],\n",
       "         [-1.0901, -0.1701,  1.8602,  ...,  4.1975, -0.5372,  0.0228],\n",
       "         ...,\n",
       "         [-0.4138, -0.1170,  2.5516,  ...,  3.9977, -0.0641, -0.0210],\n",
       "         [-0.4001, -0.0653,  2.7681,  ...,  4.3527, -0.3256,  0.2311],\n",
       "         [-1.1276, -0.0463,  2.8357,  ...,  3.5348, -0.7769,  0.0353]],\n",
       "\n",
       "        [[-1.7893, -0.6535,  3.0787,  ...,  4.9041, -0.8096,  0.3017],\n",
       "         [-0.9330, -0.9875,  1.3142,  ...,  3.6899, -0.8323, -0.8081],\n",
       "         [-0.3284, -0.4681,  1.8966,  ...,  4.0968, -0.8748,  0.2006],\n",
       "         ...,\n",
       "         [-1.0436,  0.0844,  0.9766,  ...,  4.4333, -1.5832, -0.0730],\n",
       "         [-0.7852, -0.9965,  1.6833,  ...,  3.9696, -0.4481,  0.8139],\n",
       "         [-0.6423, -1.1836,  1.0185,  ...,  4.7704, -0.2599,  0.2000]]],\n",
       "       grad_fn=<AddBackward0>)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "enc(torch.rand(5, 10, 512))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The encoder consists of 6 consecutive encoder blocks, so can simply be implemented like the following"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TransformerEncoder(nn.Module):\n",
    "    level = TensorLoggingLevels.enc_dec\n",
    "    def __init__(self, n_blocks=6, d_model=512,\n",
    "                 n_heads=8, d_ff=2048, dropout=0.1):\n",
    "        super().__init__()\n",
    "        self.encoders = nn.ModuleList([\n",
    "            EncoderBlock(d_model=d_model, d_feature=d_model // n_heads,\n",
    "                         d_ff=d_ff, dropout=dropout)\n",
    "            for _ in range(n_blocks)\n",
    "        ])\n",
    "    \n",
    "    def forward(self, x: torch.FloatTensor, mask=None):\n",
    "        for encoder in self.encoders:\n",
    "            x = encoder(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The Decoder"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The decoder is mostly the same as the encoder. There's just one additional multi-head attention block that takes the target sentence as input."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![image](https://i1.wp.com/mlexplained.com/wp-content/uploads/2017/12/%E3%82%B9%E3%82%AF%E3%83%AA%E3%83%BC%E3%83%B3%E3%82%B7%E3%83%A7%E3%83%83%E3%83%88-2017-12-29-19.14.47.png?w=287)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The keys and values are the outputs of the encoder, and the queries are the outputs of the multi-head attention over the target entence embeddings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DecoderBlock(nn.Module):\n",
    "    level = TensorLoggingLevels.enc_dec_block\n",
    "    def __init__(self, d_model=512, d_feature=64,\n",
    "                 d_ff=2048, n_heads=8, dropout=0.1):\n",
    "        super().__init__()\n",
    "        self.masked_attn_head = MultiHeadAttention(d_model, d_feature, n_heads, dropout)\n",
    "        self.attn_head = MultiHeadAttention(d_model, d_feature, n_heads, dropout)\n",
    "        self.position_wise_feed_forward = nn.Sequential(\n",
    "            nn.Linear(d_model, d_ff),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(d_ff, d_model),\n",
    "        )\n",
    "\n",
    "        self.layer_norm1 = LayerNorm(d_model)\n",
    "        self.layer_norm2 = LayerNorm(d_model)\n",
    "        self.layer_norm3 = LayerNorm(d_model)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "    def forward(self, x, enc_out, \n",
    "                src_mask=None, tgt_mask=None):\n",
    "        # Apply attention to inputs\n",
    "        att = self.masked_attn_head(x, x, x, mask=src_mask)\n",
    "        x = x + self.dropout(self.layer_norm1(att))\n",
    "        # Apply attention to the encoder outputs and outputs of the previous layer\n",
    "        att = self.attn_head(queries=x, keys=enc_out, values=enc_out, mask=tgt_mask)\n",
    "        x = x + self.dropout(self.layer_norm2(att))\n",
    "        # Apply position-wise feedforward network\n",
    "        pos = self.position_wise_feed_forward(x)\n",
    "        x = x + self.dropout(self.layer_norm2(pos))\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[EncoderBlock] Encoder block input size=torch.Size([5, 10, 512])\n",
      "[MultiHeadAttention] Input queries size=torch.Size([5, 10, 512])\n",
      "[MultiHeadAttention] output of single head size=torch.Size([5, 10, 64])\n",
      "[MultiHeadAttention] concatenated output size=torch.Size([5, 10, 512])\n",
      "[MultiHeadAttention] projected output size=torch.Size([5, 10, 512])\n",
      "[EncoderBlock] Attention output size=torch.Size([5, 10, 512])\n",
      "[EncoderBlock] Feedforward output size=torch.Size([5, 10, 512])\n",
      "[EncoderBlock] Encoder size output size=torch.Size([5, 10, 512])\n",
      "[MultiHeadAttention] Input queries size=torch.Size([5, 10, 512])\n",
      "[MultiHeadAttention] output of single head size=torch.Size([5, 10, 64])\n",
      "[MultiHeadAttention] concatenated output size=torch.Size([5, 10, 512])\n",
      "[MultiHeadAttention] projected output size=torch.Size([5, 10, 512])\n",
      "[MultiHeadAttention] Input queries size=torch.Size([5, 10, 512])\n",
      "[MultiHeadAttention] output of single head size=torch.Size([5, 10, 64])\n",
      "[MultiHeadAttention] concatenated output size=torch.Size([5, 10, 512])\n",
      "[MultiHeadAttention] projected output size=torch.Size([5, 10, 512])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[[-3.2290,  0.3235,  0.1065,  ..., -2.4907, -1.7681, -0.7010],\n",
       "         [-1.5472, -2.2766,  0.2698,  ..., -3.2174, -0.7116, -1.0080],\n",
       "         [-2.5818, -1.7114, -0.1525,  ..., -3.4856, -1.1320, -0.2742],\n",
       "         ...,\n",
       "         [-1.9990, -0.7424, -0.1879,  ..., -4.0495, -2.0523, -0.8105],\n",
       "         [-2.6684, -0.6335,  0.1777,  ..., -2.4918, -1.2974, -0.7625],\n",
       "         [-3.1428, -3.6049, -0.0199,  ..., -3.4796, -0.6485, -0.1586]],\n",
       "\n",
       "        [[-2.9487, -0.8596,  1.5827,  ..., -2.5604, -0.1855, -1.1804],\n",
       "         [-3.5555, -1.3430,  0.7323,  ..., -1.9973,  1.3033,  0.3560],\n",
       "         [-3.7100, -1.6735,  0.5658,  ..., -2.1933,  0.0994, -0.2620],\n",
       "         ...,\n",
       "         [-4.5006, -1.8862,  0.6199,  ..., -3.7571,  0.4341, -0.7318],\n",
       "         [-4.1769, -1.6887,  0.3228,  ..., -1.7997,  0.9390, -0.0960],\n",
       "         [-3.2155, -1.7711, -0.0478,  ..., -4.2723, -0.5255, -0.0749]],\n",
       "\n",
       "        [[-3.0206, -1.1389,  0.3656,  ..., -2.1171,  1.4828, -1.0802],\n",
       "         [-3.6504, -2.5347,  1.1674,  ..., -2.1860, -0.5396, -0.3272],\n",
       "         [-2.2411, -2.8799,  0.8310,  ..., -3.8143,  0.6776, -0.1047],\n",
       "         ...,\n",
       "         [-3.6718, -1.3013,  1.0194,  ..., -4.1069,  1.1408, -1.7139],\n",
       "         [-3.8910, -3.3254,  0.2053,  ..., -3.3234, -0.0180, -0.2249],\n",
       "         [-3.0260, -1.8834,  0.8524,  ..., -1.9466, -0.4605, -0.9954]],\n",
       "\n",
       "        [[-4.1134, -2.4485,  1.5797,  ..., -1.4346, -1.4182,  0.1790],\n",
       "         [-2.6200, -1.5412,  0.1274,  ..., -2.5102, -0.3610, -0.8767],\n",
       "         [-2.9746, -2.4963,  0.3496,  ..., -3.9460,  0.1921,  0.0571],\n",
       "         ...,\n",
       "         [-3.6739, -2.5789, -0.3001,  ..., -4.4960, -0.1253, -0.0794],\n",
       "         [-3.2948, -0.7941,  1.3575,  ..., -3.3816, -0.8512, -0.0411],\n",
       "         [-3.2972, -1.9540,  1.6959,  ..., -3.1695, -0.2085, -0.9750]],\n",
       "\n",
       "        [[-1.6805, -1.0441, -0.4501,  ..., -2.7119,  1.8344, -1.6279],\n",
       "         [-2.7664,  0.8360,  0.8293,  ..., -4.8047, -1.0310, -0.2812],\n",
       "         [-2.7642, -0.0724,  0.3339,  ..., -2.3763, -0.4796, -1.0766],\n",
       "         ...,\n",
       "         [-4.5325, -2.1072,  0.9622,  ..., -4.1719,  0.2678, -0.5174],\n",
       "         [-1.3087, -1.2424,  0.4613,  ..., -4.9549, -1.0824, -1.2659],\n",
       "         [-4.0965, -2.3979,  1.2412,  ..., -3.8797, -1.4461, -0.6788]]],\n",
       "       grad_fn=<AddBackward0>)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dec = DecoderBlock()\n",
    "dec(torch.rand(5, 10, 512), enc(torch.rand(5, 10, 512)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Again, the decoder is just a stack of the underlying block so is simple to implement."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TransformerDecoder(nn.Module):\n",
    "    level = TensorLoggingLevels.enc_dec\n",
    "    def __init__(self, n_blocks=6, d_model=512, d_feature=64,\n",
    "                 d_ff=2048, n_heads=8, dropout=0.1):\n",
    "        super().__init__()\n",
    "        self.position_embedding = PositionalEmbedding(d_model)\n",
    "        self.decoders = nn.ModuleList([\n",
    "            DecoderBlock(d_model=d_model, d_feature=d_model // n_heads,\n",
    "                         d_ff=d_ff, dropout=dropout)\n",
    "            for _ in range(n_blocks)\n",
    "        ])\n",
    "        \n",
    "    def forward(self, x: torch.FloatTensor, \n",
    "                enc_out: torch.FloatTensor, \n",
    "                src_mask=None, tgt_mask=None):\n",
    "        for decoder in self.decoders:\n",
    "            x = decoder(x, enc_out, src_mask=src_mask, tgt_mask=tgt_mask)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Positional Embeddings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Attention blocks are just simple matrix multiplications: therefore they don't have any notion of order! The Transformer explicitly adds positional information via the positional embeddings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PositionalEmbedding(nn.Module):\n",
    "    level = 1\n",
    "    def __init__(self, d_model, max_len=512):\n",
    "        super().__init__()        \n",
    "        # Compute the positional encodings once in log space.\n",
    "        pe = torch.zeros(max_len, d_model)\n",
    "        position = torch.arange(0, max_len).unsqueeze(1).float()\n",
    "        div_term = torch.exp(torch.arange(0, d_model, 2).float() *\n",
    "                             -(math.log(10000.0) / d_model))\n",
    "        pe[:, 0::2] = torch.sin(position * div_term)\n",
    "        pe[:, 1::2] = torch.cos(position * div_term)\n",
    "        pe = pe.unsqueeze(0)\n",
    "        self.weight = nn.Parameter(pe, requires_grad=False)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        return self.weight[:, :x.size(1), :] # (1, Seq, Feature)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "class WordPositionEmbedding(nn.Module):\n",
    "    level = 1\n",
    "    def __init__(self, vocab_size, d_model=512):\n",
    "        super().__init__()\n",
    "        self.word_embedding = nn.Embedding(vocab_size, d_model)\n",
    "        self.position_embedding = PositionalEmbedding(d_model)\n",
    "        \n",
    "    def forward(self, x: torch.LongTensor, mask=None) -> torch.FloatTensor:\n",
    "        return self.word_embedding(x) + self.position_embedding(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "emb = WordPositionEmbedding(1000)\n",
    "encoder = TransformerEncoder()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[EncoderBlock] Encoder block input size=torch.Size([5, 30, 512])\n",
      "[MultiHeadAttention] Input queries size=torch.Size([5, 30, 512])\n",
      "[MultiHeadAttention] output of single head size=torch.Size([5, 30, 64])\n",
      "[MultiHeadAttention] concatenated output size=torch.Size([5, 30, 512])\n",
      "[MultiHeadAttention] projected output size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Attention output size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Feedforward output size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Encoder size output size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Encoder block input size=torch.Size([5, 30, 512])\n",
      "[MultiHeadAttention] Input queries size=torch.Size([5, 30, 512])\n",
      "[MultiHeadAttention] output of single head size=torch.Size([5, 30, 64])\n",
      "[MultiHeadAttention] concatenated output size=torch.Size([5, 30, 512])\n",
      "[MultiHeadAttention] projected output size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Attention output size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Feedforward output size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Encoder size output size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Encoder block input size=torch.Size([5, 30, 512])\n",
      "[MultiHeadAttention] Input queries size=torch.Size([5, 30, 512])\n",
      "[MultiHeadAttention] output of single head size=torch.Size([5, 30, 64])\n",
      "[MultiHeadAttention] concatenated output size=torch.Size([5, 30, 512])\n",
      "[MultiHeadAttention] projected output size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Attention output size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Feedforward output size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Encoder size output size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Encoder block input size=torch.Size([5, 30, 512])\n",
      "[MultiHeadAttention] Input queries size=torch.Size([5, 30, 512])\n",
      "[MultiHeadAttention] output of single head size=torch.Size([5, 30, 64])\n",
      "[MultiHeadAttention] concatenated output size=torch.Size([5, 30, 512])\n",
      "[MultiHeadAttention] projected output size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Attention output size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Feedforward output size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Encoder size output size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Encoder block input size=torch.Size([5, 30, 512])\n",
      "[MultiHeadAttention] Input queries size=torch.Size([5, 30, 512])\n",
      "[MultiHeadAttention] output of single head size=torch.Size([5, 30, 64])\n",
      "[MultiHeadAttention] concatenated output size=torch.Size([5, 30, 512])\n",
      "[MultiHeadAttention] projected output size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Attention output size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Feedforward output size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Encoder size output size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Encoder block input size=torch.Size([5, 30, 512])\n",
      "[MultiHeadAttention] Input queries size=torch.Size([5, 30, 512])\n",
      "[MultiHeadAttention] output of single head size=torch.Size([5, 30, 64])\n",
      "[MultiHeadAttention] concatenated output size=torch.Size([5, 30, 512])\n",
      "[MultiHeadAttention] projected output size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Attention output size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Feedforward output size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Encoder size output size=torch.Size([5, 30, 512])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[[  0.3980,  -4.3829,  -7.4285,  ...,   3.1774,   1.7740,   2.5976],\n",
       "         [  0.0647,  -4.1780,  -7.2751,  ...,   3.9347,  -0.8255,   5.2059],\n",
       "         [  0.3450,  -5.0790,  -6.8530,  ...,   1.7033,   2.2193,   0.8902],\n",
       "         ...,\n",
       "         [  4.6940,  -4.9993, -11.5730,  ...,   1.8663,   1.8704,   1.3255],\n",
       "         [  0.1014,  -5.3799,  -7.3519,  ...,   1.3983,   3.2914,   3.3760],\n",
       "         [  0.2468,  -3.6646,  -9.4579,  ...,   5.5852,   1.0168,   1.5736]],\n",
       "\n",
       "        [[  5.9047,  -4.4927,  -8.4207,  ...,   2.1502,   3.5834,   4.6961],\n",
       "         [  4.1149,   0.2121,  -3.7838,  ...,   3.4321,  -0.5452,   2.7152],\n",
       "         [  3.0612,  -4.8328,  -4.9506,  ...,   1.6598,   5.2516,   2.9424],\n",
       "         ...,\n",
       "         [  4.9755,  -3.8555,  -7.2134,  ...,   4.7120,   5.0917,   3.8296],\n",
       "         [  9.1925,  -4.7523,  -4.5528,  ...,   2.6217,   3.4018,  -1.0523],\n",
       "         [  5.9289,  -4.7409,  -7.0705,  ...,   4.5752,   2.3157,   1.0976]],\n",
       "\n",
       "        [[  5.5719,  -3.3389,  -8.9670,  ...,   3.2728,   7.8318,   4.9943],\n",
       "         [  1.8820,  -4.5469,  -4.7731,  ...,   6.5021,   1.4791,   0.9152],\n",
       "         [  0.4770,  -6.0725,  -8.8598,  ...,   4.8887,   2.5641,   0.2976],\n",
       "         ...,\n",
       "         [  5.9999,  -4.4015,  -9.4782,  ...,   2.3594,   5.5085,   2.6492],\n",
       "         [  5.0709,  -3.2256,  -2.6433,  ...,   2.4225,   3.5039,   7.1794],\n",
       "         [  4.1746,  -7.4886, -13.7985,  ...,   6.3600,   4.6518,   2.6611]],\n",
       "\n",
       "        [[  1.5971,  -0.9796,  -8.7406,  ...,   1.4917,   7.1294,  -0.5388],\n",
       "         [  8.2672,   2.8524,  -7.3631,  ...,   3.6907,   2.7596,  -0.0145],\n",
       "         [  8.4459,  -4.4426,  -6.8800,  ...,   2.4980,   3.5902,   1.2568],\n",
       "         ...,\n",
       "         [  4.7770,  -1.5967,  -4.4558,  ...,   2.9476,   8.6169,   2.4252],\n",
       "         [  1.8154,  -4.8706,  -5.9819,  ...,   5.6543,   5.5660,   3.7610],\n",
       "         [  3.8060,  -4.0780,  -9.1496,  ...,   3.6763,  -0.4940,  -3.0517]],\n",
       "\n",
       "        [[  3.6584,  -4.0857,  -7.9289,  ...,  -1.7100,   2.8362,   3.4344],\n",
       "         [  0.7226,  -5.0240,  -3.9058,  ...,  -0.4559,   7.2419,  -0.4174],\n",
       "         [  2.6257,  -3.5737,  -8.4171,  ...,   0.5259,   2.1638,  -2.0950],\n",
       "         ...,\n",
       "         [  4.2735,  -8.0241,  -9.0453,  ...,   3.6112,   5.8559,   6.0040],\n",
       "         [ -0.0731,  -3.0447,  -9.2193,  ...,   1.4438,   2.9617,   1.0221],\n",
       "         [  2.3344,  -0.9475,  -7.1632,  ...,   5.6997,   4.1256,  -2.8659]]],\n",
       "       grad_fn=<AddBackward0>)"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "encoder(emb(torch.randint(1000, (5, 30))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Putting it All Together"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's put everything together now."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![image](https://camo.githubusercontent.com/88e8f36ce61dedfd2491885b8df2f68c4d1f92f5/687474703a2f2f696d6775722e636f6d2f316b72463252362e706e67)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We'll supress logging from the scaled dot product attention now\n",
    "logger.setLevel(TensorLoggingLevels.enc_dec_block)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "emb = WordPositionEmbedding(1000)\n",
    "encoder = TransformerEncoder()\n",
    "decoder = TransformerDecoder()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[EncoderBlock] Encoder block input size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Attention output size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Feedforward output size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Encoder size output size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Encoder block input size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Attention output size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Feedforward output size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Encoder size output size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Encoder block input size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Attention output size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Feedforward output size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Encoder size output size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Encoder block input size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Attention output size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Feedforward output size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Encoder size output size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Encoder block input size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Attention output size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Feedforward output size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Encoder size output size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Encoder block input size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Attention output size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Feedforward output size=torch.Size([5, 30, 512])\n",
      "[EncoderBlock] Encoder size output size=torch.Size([5, 30, 512])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[[-3.9986,  5.7489,  6.7348,  ..., -6.9215,  7.3174,  3.6647],\n",
       "         [ 1.9130, -1.8961,  3.4541,  ..., -6.9436,  4.5712,  3.4783],\n",
       "         [-3.4003,  2.4626,  4.9213,  ..., -5.5661,  3.2948,  4.1336],\n",
       "         ...,\n",
       "         [-1.5915, -0.5826,  7.6223,  ...,  0.1395,  5.7651,  3.3921],\n",
       "         [-3.5197, -3.5261,  5.5072,  ..., -6.2547,  1.4598,  3.8519],\n",
       "         [-2.9156, -0.7501,  3.9160,  ..., -5.9510,  2.6685,  1.9906]],\n",
       "\n",
       "        [[-0.7690, -1.0143,  1.3752,  ..., -6.4549,  2.9288,  4.6994],\n",
       "         [-3.7906, -3.1318,  1.3130,  ..., -8.5190,  2.9146,  3.7494],\n",
       "         [-0.5046, -2.0664,  5.1320,  ..., -7.0491,  2.5344,  4.0302],\n",
       "         ...,\n",
       "         [-1.2325, -1.2299,  1.2333,  ..., -6.6735,  1.6056,  4.2359],\n",
       "         [-2.7019, -2.6522, -1.4139,  ..., -4.7800,  1.5574,  2.3967],\n",
       "         [ 0.3484, -4.1592,  2.2151,  ..., -6.2922,  1.4931,  0.0246]],\n",
       "\n",
       "        [[-0.9702,  3.6488,  5.2317,  ..., -5.5706,  2.4317,  4.3568],\n",
       "         [-5.4813, -0.7727,  7.8615,  ..., -5.8813,  3.5327,  5.1189],\n",
       "         [-2.0725, -4.1126,  5.5792,  ..., -4.4119,  6.0067,  5.1591],\n",
       "         ...,\n",
       "         [-2.4875, -0.9064,  2.5081,  ..., -2.3326, -1.6754,  0.2841],\n",
       "         [-4.1770, -2.9200,  0.6833,  ..., -4.5765,  2.3240,  6.7224],\n",
       "         [-4.4732, -1.7839,  1.8468,  ..., -3.8334,  0.8775,  6.5711]],\n",
       "\n",
       "        [[-4.6256, -2.4793,  2.5954,  ..., -5.1335, -0.2744,  2.7417],\n",
       "         [-3.2865, -0.2110,  4.1452,  ..., -4.3811,  2.1002,  1.0985],\n",
       "         [-2.2448, -1.2541,  5.9155,  ..., -2.8403,  0.9341,  1.6664],\n",
       "         ...,\n",
       "         [-2.5197, -2.7198,  6.9956,  ..., -7.5234,  0.8126,  4.4213],\n",
       "         [-2.9402, -4.0223,  1.5813,  ..., -0.5574,  2.9243, -1.2572],\n",
       "         [-3.4166, -4.3764,  1.1802,  ..., -3.3221,  0.2821,  6.5111]],\n",
       "\n",
       "        [[-2.5195, -0.4850,  7.5954,  ..., -2.2857,  0.7949,  1.4315],\n",
       "         [-2.8258, -2.9534,  3.6664,  ..., -3.0961,  2.3875,  5.3663],\n",
       "         [ 0.7049, -4.9491,  6.6818,  ..., -7.6540,  2.5320,  4.9328],\n",
       "         ...,\n",
       "         [ 1.0463, -1.3215,  7.6014,  ..., -4.5010, -1.6455,  2.8965],\n",
       "         [ 0.7319, -8.3325,  3.7722,  ..., -3.4189,  1.3441,  4.2740],\n",
       "         [-0.1283, -2.4150,  2.1482,  ..., -4.2556,  1.2888,  4.7413]]],\n",
       "       grad_fn=<AddBackward0>)"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "src_ids = torch.randint(1000, (5, 30))\n",
    "tgt_ids = torch.randint(1000, (5, 30))\n",
    "x = encoder(emb(src_ids))\n",
    "decoder(emb(tgt_ids), x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
